"""
All the things related to heatmap evaluation
"""
import numpy as np


# Container for heatmap evaluation results
class HeatmapResult(object):
    def __init__(self, num_intervals=1000):
        # record the data
        # _precision and _recall accumulates counts in different ways (due to the distance tolerance).
        self.data = {
            "tp_precision": np.zeros([num_intervals, 1]),
            "fp_precision": np.zeros([num_intervals, 1]),
            "tp_recall": np.zeros([num_intervals, 1]),
            "num_points": 0
        }

    # Accumulate the results
    def update(self, heatmap_res):
        # Update the accumulative counts
        self.data["tp_precision"] += heatmap_res["tp_precision"]
        self.data["fp_precision"] += heatmap_res["fp_precision"]
        self.data["tp_recall"] += heatmap_res["tp_recall"]
        self.data["num_points"] += heatmap_res["num_gt"]

    # Compute the precision and recall
    def compute_pr(self):
        tp_precision = np.squeeze(self.data["tp_precision"], axis=-1)
        fp_precision = np.squeeze(self.data["fp_precision"], axis=-1)
        tp_recall = np.squeeze(self.data["tp_recall"], axis=-1)

        precision = self.div0(tp_precision, tp_precision + fp_precision)
        recall = self.div0(tp_recall, self.data["num_points"])

        # Add entries for plotting
        recall = np.concatenate([[0], recall, [1]])
        precision = np.concatenate([[0], precision, [0]])
        # TODO: Get the max accumulated results ??
        precision = np.maximum.accumulate(precision[::-1])[::-1]

        # Compute mAP
        mAP = self.compute_mAP(precision, recall)

        return {
            "precision": precision,
            "recall": recall,
            "mAP": mAP
        }

    # Compute mAP
    @staticmethod
    def compute_mAP(precision, recall):
        return np.sum(precision[1:] * (recall[1:] - recall[:-1]))

    @staticmethod
    def div0(a, b):
        with np.errstate(divide='ignore', invalid='ignore'):
            c = np.true_divide(a, b)
            idx = ~np.isfinite(c)
            c[idx] = np.where(a[idx] == 0, 1, 0)  # -inf inf NaN
        return c


# Compute tp fp for heatmap prediction (using distance metric).
# TODO: extend to batch size > 1
def compute_tp_fp(heatmap_pred, heatmap_gt, valid_mask, dist_thresh_ratio=0.01, num_intervals=1000):
    # heatmap_pred H x W x 2
    # heatmap_gt H x W x 1
    # valid_mask H x W x 1
    heatmap_pred = heatmap_pred[..., 1]

    # Compute the distance threshold based on the length of image diagonal
    h, w = heatmap_pred.shape
    dist_thresh = round(((h ** 2 + w ** 2) ** 0.5) * dist_thresh_ratio)

    # Compute the gt location
    # 元组类型，分量为2，分别保存 H W
    gt_loc = np.where(heatmap_gt[..., 0] * valid_mask[..., 0])
    # Convert into Nx2 format.
    gt_loc = np.concatenate([gt_loc[0][..., None], gt_loc[1][..., None]], axis=-1)
    num_gt = gt_loc.shape[0]

    # Process the prediction data
    heatmap_pred_valid = heatmap_pred * valid_mask[..., 0]

    # Iterate through all the prob thresholds
    tp_precision_array = np.zeros([num_intervals, 1])
    fp_precision_array = np.zeros([num_intervals, 1])
    tp_recall_array = np.zeros([num_intervals, 1])
    thresh_intervals = np.linspace(0.95, 0.05, num_intervals, endpoint=True)

    for idx, thresh in enumerate(thresh_intervals):
        # print("debug", idx)
        if idx == 0:
            thresh_prev = 1.
        else:
            thresh_prev = thresh_intervals[idx - 1]
        # H x W 的布尔型数组
        mask1 = heatmap_pred_valid > thresh
        mask2 = heatmap_pred_valid <= thresh_prev
        mask = np.where(mask1 * mask2)

        # Convert prediction to M*2 format
        pred = np.concatenate([mask[0][..., None], mask[1][..., None]], axis=-1)

        # Compute precision related statistics
        # M x 1 x 2 - 1 x N x 2
        diff_precision = np.expand_dims(pred, axis=1) - gt_loc[None, ...]
        # M x N
        dist_precision = np.linalg.norm(diff_precision, axis=-1)
        # M x N 的布尔型数组
        matches_precision = np.less_equal(dist_precision, dist_thresh)

        tp_precision = np.any(matches_precision, axis=1)
        fp_precision = np.logical_not(tp_precision)

        # Compute the recall related statistics
        # N x 1 x 2 - 1 x M x 2
        diff_recall = np.expand_dims(gt_loc, axis=1) - pred[None, ...]
        # N x M
        dist_recall = np.linalg.norm(diff_recall, axis=-1)

        # Filter out overlap counts
        if idx == 0:
            # N x M
            matches_recall = np.less_equal(dist_recall, dist_thresh)
            tp_recall = np.any(matches_recall, axis=1)
        else:
            matches_recall = np.less_equal(dist_recall, dist_thresh)
            tp_recall = np.any(matches_recall, axis=1)
            # 逻辑或
            tp_recall = np.logical_or(tp_recall, tp_recall_prev)

        tp_recall_prev = tp_recall

        # Record the results for single threshold
        tp_precision_array[idx] = np.sum(tp_precision.astype(np.int))
        fp_precision_array[idx] = np.sum(fp_precision.astype(np.int))
        tp_recall_array[idx] = np.sum(tp_recall.astype(np.int))

    # Compute the cumulative sum
    tp_precision_array = np.cumsum(tp_precision_array)[..., None]
    fp_precision_array = np.cumsum(fp_precision_array)[..., None]

    return {
        "tp_precision": tp_precision_array,
        "fp_precision": fp_precision_array,
        "tp_recall": tp_recall_array,
        "num_gt": num_gt
    }
